#include "irtest.h"

#include <sys/types.h>

#include <cstdio>
#include <iostream>
#include <ostream>
#include <string>

#include "arithmetic.h"
#include "irfunction.h"
#include "semantics.h"
#include "types.h"
#include "utils.h"

ArithValue& IRTest::GetValue(IR::Operand* x) {
    Utils::Assert(x, "");
    if (auto imm = x->As<IR::Immediate>()) {
        return imm->As<IR::Immediate>()->value;
    }
    auto v = x->As<IR::Variable>();
    if (v->isGlobal) return global[v->name];
    if (!var.count(v->name)) {
        if (v->type.base == Type::Basic::Float) {
            var[v->name] = ArithValue(.0f);
        } else if (v->type.base == Type::Basic::Int) {
            var[v->name] = ArithValue(0);
        }
    }
    return var[v->name];
}

IRBlock* IRTest::Step(IR::IRBase* irBase) {
    if (auto ir = irBase->As<IR::Operation>()) {
        auto r = GetValue(ir->rhs);
        auto& dest = GetValue(ir->dest);
        if (ir->lhs) {
            auto l = GetValue(ir->lhs);
            dest = ArithOp::Calc(l, r, ir->op);
        } else {
            dest = ArithOp::Calc(r, ir->op);
        }
        return nullptr;
    }
    if (auto ir = irBase->As<IR::Alloc>()) {
        GetValue(ir->dest) = ArithValue(memPtr);
        for (int i = 0; i < ir->size; ++i) {
            mem[memPtr + i] = ArithValue(0);
        }
        memPtr += ir->size;
        return nullptr;
    }
    if (auto ir = irBase->As<IR::Load>()) {
        int base = GetValue(ir->base).v.i;
        int off = GetValue(ir->offset).v.i;
        GetValue(ir->dest) = mem[base + off];
        return nullptr;
    }
    if (auto ir = irBase->As<IR::Store>()) {
        int base = GetValue(ir->base).v.i;
        int off = GetValue(ir->offset).v.i;
        mem[base + off] = GetValue(ir->src);
        return nullptr;
    }
    if (auto ir = irBase->As<IR::Call>()) {
        std::vector<ArithValue> args;
        for (auto arg : ir->args) args.push_back(GetValue(arg));
        auto backup = var;
        var.clear();
        auto ret = Call(ir->func, args);
        var = backup;
        if (ir->dest) GetValue(ir->dest) = ret;
        return nullptr;
    }
    if (auto ir = irBase->As<IR::Return>()) {
        return nullptr;
    }
    if (auto ir = irBase->As<IR::Jump>()) {
        return ir->nxt;
    }
    if (auto ir = irBase->As<IR::Branch>()) {
        auto l = GetValue(ir->lhs);
        auto r = GetValue(ir->rhs);
        auto cond = ArithOp::Calc(l, r, ir->op);
        if (cond.type == Type::Basic::Int) {
            return cond.v.i ? ir->t : ir->f;
        } else {
            return cond.v.f ? ir->t : ir->f;
        }
    }
    return nullptr;
}

ArithValue IRTest::Call(IRFunction* func, const std::vector<ArithValue>& args) {
    if (!func->entry) {
        return CallLib(func->name, args);
    }

    Utils::Assert(func->args.size() == args.size(), "argument size error");
    for (int i = 0; i < args.size(); ++i) {
        GetValue(func->args[i]) = args[i];
    }

    auto block = func->entry;
    while (block != func->exit) {
        IRBlock* nxt;
        for (auto ir : block->irList) {
            nxt = Step(ir);
        }
        for (auto phi : nxt->phis) {
            GetValue(phi->dest) = GetValue(phi->srcs[block]);
        }
        block = nxt;
        Utils::Assert(nxt, "");
    }
    for (auto ir : block->irList) {
        if (auto ret = ir->As<IR::Return>()) {
            if (!ret->ret) return ArithValue(0);
            return GetValue(ret->ret);
        }
        Step(ir);
    }
    return ArithValue(0);
}

void IRTest::Init() {
    for (auto [v, sym] : module->globalVars) {
        if (auto info = sym->info->As<Semantics::LValue>()) {
            global[v->name] = info->init;
        } else if (auto info = sym->info->As<Semantics::Pointer>()) {
            global[v->name] = ArithValue(memPtr);
            for (int i = memPtr; i < memPtr + sym->allocMem; ++i) {
                mem[i] = ArithValue(info->baseType);
            }
            if (info->init) {
                for (auto [x, val] : *info->init) {
                    mem[memPtr + x] = val;
                }
            }
        }
        memPtr += sym->allocMem;
    }
}

void* IRTest::Run(void* input) {
    auto module = (IRModule*)input;
    this->module = module;
    Utils::Assert(module->funcs.count(entry), "entry function not exists");
    std::cout << RED "\n[IRTest]\n" DEF;
    Init();
    ofs << Call(module->funcs[entry], entryArgs).v.i << std::endl;
    return module;
}

ArithValue IRTest::CallLib(const std::string& name, const ArgList& args) {
    auto getInt = [&]() {
        int t;
        ifs >> t;
        return ArithValue(t);
    };
    auto putInt = [&](ArithValue i) { ofs << i.v.i; };
    auto getChar = [&]() { return ArithValue(ifs.get()); };
    auto putChar = [&](ArithValue c) { ofs << char(c.v.i); };
    auto getFloat = [&]() {
        float t;
        std::string s;
        ifs >> s;
        sscanf(s.c_str(), "%a", &t);
        return ArithValue(t);
    };
    auto putFloat = [&](ArithValue f) { ofs << Utils::Format("%a", f.v.f); };

    std::unordered_map<std::string, LibFunc> funcMap;

    funcMap["getint"] = [&](const ArgList& args) { return getInt(); };
    funcMap["getch"] = [&](const ArgList& args) { return getChar(); };
    funcMap["getfloat"] = [&](const ArgList& args) { return getFloat(); };
    funcMap["getarray"] = [&](const ArgList& args) {
        int a = args[0].v.i;
        auto n = getInt();
        for (int i = 0; i < n.v.i; ++i) {
            mem[a + i] = getInt();
        }
        return n;
    };
    funcMap["getfarray"] = [&](const ArgList& args) {
        int a = args[0].v.i;
        auto n = getInt();
        for (int i = 0; i < n.v.i; ++i) {
            mem[a + i] = getFloat();
        }
        return n;
    };

    funcMap["putint"] = [&](const ArgList& args) {
        putInt(args[0]);
        return ArithValue(0);
    };
    funcMap["putch"] = [&](const ArgList& args) {
        putChar(args[0]);
        return ArithValue(0);
    };
    funcMap["putfloat"] = [&](const ArgList& args) {
        putFloat(args[0]);
        return ArithValue(0);
    };
    funcMap["putarray"] = [&](const ArgList& args) {
        int n = args[0].v.i;
        int a = args[1].v.i;
        ofs << n << ":";
        for (int i = 0; i < n; ++i) {
            ofs << " ";
            putInt(mem[a + i]);
        }
        ofs << "\n";
        return ArithValue(0);
    };
    funcMap["putfarray"] = [&](const ArgList& args) {
        int n = args[0].v.i;
        int a = args[1].v.i;
        ofs << n << ":";
        for (int i = 0; i < n; ++i) {
            ofs << " ";
            putFloat(mem[a + i]);
        }
        ofs << "\n";
        return ArithValue(0);
    };

    funcMap["putf"] = [&](const ArgList& args) {
        int a = args[0].v.i;
        int pos = 1;
        std::string fmt, ret;
        for (int i = 0;; ++i) {
            char c = ((char*)&mem[a + i / 4].v.i)[i % 4];
            fmt += c;
            if (!c) break;
        }
        for (auto it = fmt.begin(); it != fmt.end(); ++it) {
            if (*it != '%') {
                ret += *it;
                continue;
            }
            ++it;
            switch (*it) {
                case '%':
                    ret += '%';
                    break;
                case 'd':
                    ret += std::to_string(args[pos++].v.i);
                    break;
                case 'c':
                    ret += char(args[pos++].v.i);
                    break;
                case 'f':
                    ret += std::to_string(args[pos++].v.f);
                    break;
            }
        }
        ofs << ret;
        return ArithValue(0);
    };

    if (funcMap.count(name)) return funcMap[name](args);
    return ArithValue(0);
}